// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//


#pragma once

#include <pybind11/embed.h>
#include <pollux/type/type.h>

namespace kumo::pollux::py {

class PyType {
 public:
  explicit PyType(const TypePtr& type = nullptr) : type_(type) {}

  std::string toString() const {
    if (!type_) {
      return "[nullptr]";
    }
    return type_->toString();
  }

  TypePtr type() const {
    return type_;
  }

  bool equivalent(const PyType& other) const {
    return type_->equivalent(*other.type());
  }

  // Factory functions:

  static PyType createRowType(
      const std::vector<std::string>& names,
      const std::vector<PyType>& pyTypes) {
    std::vector<TypePtr> types;
    for (const auto& pyType : pyTypes) {
      types.emplace_back(pyType.type());
    }

    if (names.empty()) {
      return PyType{ROW(std::move(types))};
    }
    return PyType{ROW(melon::copy(names), std::move(types))};
  }

  static PyType createMapType(const PyType& keyType, const PyType& valueType) {
    return PyType{MAP(keyType.type(), valueType.type())};
  }

  static PyType createArrayType(const PyType& elementsType) {
    return PyType{ARRAY(elementsType.type())};
  }

  static PyType createBigint() {
    return PyType{BIGINT()};
  }

  static PyType createInteger() {
    return PyType{INTEGER()};
  }

  static PyType createSmallint() {
    return PyType{SMALLINT()};
  }

  static PyType createTinyint() {
    return PyType{TINYINT()};
  }

  static PyType createBoolean() {
    return PyType{BOOLEAN()};
  }

  static PyType createReal() {
    return PyType{REAL()};
  }

  static PyType createDouble() {
    return PyType{DOUBLE()};
  }

  static PyType createDate() {
    return PyType{DATE()};
  }

  static PyType createVarchar() {
    return PyType{VARCHAR()};
  }

  static PyType createVarbinary() {
    return PyType{VARBINARY()};
  }

 private:
  TypePtr type_;
};

} // namespace kumo::pollux::py
